import torch
import pandas
import torch.nn as nn
import numpy as np
import torch.nn.init as init
import torch.nn.functional as F
from torch import nn, optim
import random
import os
from datetime import datetime
import pickle
import matplotlib.pyplot as plt
import model
import seed
seed.seed_everything()
import focal


def initialize_weights(a):
    for name, module in a.named_children():
        if isinstance(module, nn.Conv1d) or isinstance(module, nn.Linear):
            # 获取模块中的权重参数
            weight = getattr(module, 'weight')
            # 使用 He 初始化
            init.kaiming_normal_(weight, mode='fan_out', nonlinearity='relu')
            # 如果有偏置参数，也初始化为零
            if hasattr(module, 'bias') and module.bias is not None:
                bias = getattr(module, 'bias')
                init.constant_(bias, 0)
        elif isinstance(module, nn.Module):
            # 如果是子模块，递归调用 initialize_weights 方法
            initialize_weights(module)




model = model.Model()
model=model.cuda()
#initialize_weights(model)
#arcloss=torch.tensor([1.61,4.17,7.11])
#m_loss = focal.Focal_Loss(arcloss)
m_loss = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.Adam(model.parameters())

if __name__ == "__main__":
    import memorymonitor
    memorymonitor.start()
    b = pandas.read_csv(r"D:\old\Desktop\old\animal\barking-emotion-recognition\data\dataset_2.csv")
    with open(r"D:/emotiondataset/"+"trainlist","rb") as f:
        trainlist=pickle.load(f)
    loss_0=99999999999
    train_batch=6
    test_batch=8
    arc_0=0
    while True:
        model.train()
        loss=0
        sindex=[]
        for i in range(train_batch):
            index = random.choices(trainlist[random.randint(0, 2)])[0]
            filestr = r"D:/old/Desktop/old/animal/barking-emotion-recognition/data/audioset_audios/" + b["ytid"][
                index] + "_" + str(int(b["start"][index])) + "_" + str(int(b["stop"][index])) + "_cut.mp3"
            if b["label"][index] == 'Happy':
                label = torch.tensor([1, 0, 0]).float().cuda()
            elif b["label"][index] == 'Aggressive':
                label = torch.tensor([0, 1, 0]).float().cuda()
            else:
                label = torch.tensor([0, 0, 1]).float().cuda()
            try:
                s = model(filestr)
            except FileNotFoundError:
                continue
            sindex.append(s)
            # sindex.append(int(torch.argmax(s)))
            #loss = loss + m_loss(s, label)
            loss = loss + m_loss(label, torch.squeeze(s, 0))
        loss=loss/len(sindex)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        model.eval()
        test_loss=0
        test_arc=0
        test_index=[]
        with torch.no_grad():
            hi=0
            ai=0
            ci=0
            hi0=0
            ai0=0
            ci0=0
            for i in range(test_batch):
                index = random.randint(401, 604)
                filestr = r"D:/old/Desktop/old/animal/barking-emotion-recognition/data/audioset_audios/" + b["ytid"][
                    index] + "_" + str(int(b["start"][index])) + "_" + str(int(b["stop"][index])) + "_cut.mp3"
                if b["label"][index] == 'Happy':
                    label = torch.tensor([1, 0, 0]).float().cuda()
                    hi=hi+1
                elif b["label"][index] == 'Aggressive':
                    label = torch.tensor([0, 1, 0]).float().cuda()
                    ai=ai+1
                else:
                    label = torch.tensor([0, 0, 1]).float().cuda()
                    ci=ci+1
                try:
                    s = model(filestr)
                except FileNotFoundError:
                    continue
                # 交叉熵代价函数out（batch，C：类别的数量），labels（batch）
                if torch.argmax(s) == torch.argmax(label):
                    test_arc = test_arc + 1
                    if torch.argmax(s) == 0:
                        hi0 = hi0 + 1
                    if torch.argmax(s) == 1:
                        ai0 = ai0 + 1
                    if torch.argmax(s) == 2:
                        ci0 = ci0 + 1
                test_index.append(int(torch.argmax(s)))
                #test_loss = test_loss + m_loss(s, label)

                test_loss = test_loss + m_loss(label, torch.squeeze(s, 0))
        test_loss=test_loss/len(test_index)
        if test_arc/test_batch>arc_0:
            arc_0=test_arc/test_batch
            try:
                torch.save(model, 'save.pkl')
            except:
                pass
        try:
            hiarc=hi0/hi
        except ZeroDivisionError:
            hiarc="nan"
        try:
            aiarc = ai0 / ai
        except ZeroDivisionError:
            aiarc = "nan"
        try:
            ciarc = ci0 / ci
        except ZeroDivisionError:
            ciarc = "nan"
        print(float(loss),float(test_loss),test_index,test_arc/test_batch,hiarc,aiarc,ciarc)
        #print(loss,test_loss,test_arc/test_batch,optimizer.state_dict()['param_groups'][0]['lr'])
        #1/0